"""
YOLO11 推理引擎
使用 Ultralytics YOLO 官方接口（简化版）
参考: https://docs.ultralytics.com/zh/usage/python/

示例:
    from ultralytics import YOLO
    model = YOLO("/Users/yanlp/models/yolo11n.pt")
    results = model("image.jpg")
"""
import logging
from typing import Dict
import base64
import io
import tempfile
import os
from PIL import Image

logger = logging.getLogger(__name__)


def parse_yolo_results(results) -> Dict:
    """
    解析YOLO推理结果为标准格式
    
    Args:
        results: YOLO模型的推理结果
    
    Returns:
        包含检测结果的字典（Ultralytics HUB API格式）
    """
    # 解析结果（取第一个结果）
    result = results[0]
    
    # 提取检测结果
    detections = []
    if result.boxes is not None:
        boxes = result.boxes
        for i in range(len(boxes)):
            # 边界框坐标 (x1, y1, x2, y2)
            box = boxes.xyxy[i].cpu().numpy().tolist()
            # 置信度
            confidence = float(boxes.conf[i].cpu().numpy())
            # 类别
            class_id = int(boxes.cls[i].cpu().numpy())
            class_name = result.names[class_id]
            
            detection = {
                "class": class_id,
                "name": class_name,
                "confidence": confidence,
                "box": {
                    "x1": box[0],
                    "y1": box[1],
                    "x2": box[2],
                    "y2": box[3]
                }
            }
            
            # 如果有分割掩码
            if hasattr(result, 'masks') and result.masks is not None:
                mask = result.masks.xy[i]
                detection["segments"] = {
                    "x": mask[:, 0].tolist(),
                    "y": mask[:, 1].tolist()
                }
            
            # 如果有关键点（姿态估计）
            if hasattr(result, 'keypoints') and result.keypoints is not None:
                kpts = result.keypoints.xy[i].cpu().numpy()
                conf_kpts = result.keypoints.conf[i].cpu().numpy() if result.keypoints.conf is not None else None
                detection["keypoints"] = {
                    "x": kpts[:, 0].tolist(),
                    "y": kpts[:, 1].tolist(),
                    "visible": conf_kpts.tolist() if conf_kpts is not None else None
                }
            
            # 如果有OBB（旋转边界框）
            if hasattr(result, 'obb') and result.obb is not None:
                obb = result.obb.xyxyxyxy[i].cpu().numpy().tolist()
                detection["obb"] = obb
            
            detections.append(detection)
    
    logger.info(f"检测完成: {len(detections)} 个目标")
    
    # 获取版本信息
    import sys
    import torch
    import ultralytics
    
    # 返回统一格式（参考Ultralytics HUB API）
    return {
        "images": [
            {
                "results": detections,
                "shape": [result.orig_shape[0], result.orig_shape[1]],
                "speed": {
                    "inference": result.speed.get('inference', 0),
                    "preprocess": result.speed.get('preprocess', 0),
                    "postprocess": result.speed.get('postprocess', 0)
                }
            }
        ],
        "metadata": {
            "imageCount": 1,
            "version": {
                "python": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
                #"torch": torch.__version__,
                "ultralytics": ultralytics.__version__
            }
        }
    }


def predict_from_base64(model_path: str, image_base64: str, **kwargs) -> Dict:
    """
    对Base64编码的图像进行推理（动态加载模型）
    
    Args:
        model_path: YOLO模型文件路径
        image_base64: Base64编码的图像数据
        **kwargs: 推理参数（imgsz, conf, iou等）
    
    Returns:
        包含检测结果的字典
    """
    from ultralytics import YOLO
    
    temp_file = None
    try:
        # 解码Base64图像
        image_data = base64.b64decode(image_base64)
        image = Image.open(io.BytesIO(image_data))
        
        # 保存到临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as f:
            image.save(f, format='JPEG')
            temp_file = f.name
        
        logger.info(f"图像保存到临时文件: {temp_file}")
        
        # 动态加载模型并推理
        logger.info(f"加载模型: {model_path}")
        model = YOLO(model_path)
        
        # 执行推理（官方简化接口）
        results = model(temp_file, **kwargs)
        
        # 解析结果
        return parse_yolo_results(results)
        
    except Exception as e:
        logger.error(f"Base64推理失败: {str(e)}", exc_info=True)
        raise
    finally:
        # 清理临时文件
        if temp_file and os.path.exists(temp_file):
            try:
                os.unlink(temp_file)
                logger.debug(f"已删除临时文件: {temp_file}")
            except Exception as e:
                logger.warning(f"删除临时文件失败: {e}")
